import torch
import json

stride = 1
padding = 0
kernalSize = 7
imgSize = 32
inChanel = 3
outChanel = 16

input = torch.randn(inChanel, imgSize, imgSize)
weight = torch.randn(outChanel, inChanel, kernalSize, kernalSize)
bias = torch.randn(outChanel)
params = {'weight': weight, 'bias': bias}
operator = torch.nn.Conv2d(
    inChanel, outChanel, kernalSize, stride=stride, padding=padding)
operator.load_state_dict(params)
output = operator.forward(input)

exportData = [
    ('input', input),
    ('weight', weight),
    ('bias', bias),
    ('output', output)
]

data = []
for name, tensor in exportData:
    data.append({"layer": name, "value": tensor.tolist()})

data.append({'layer': 'stride', 'value': stride})
data.append({'layer': 'padding', 'value': padding})

jsonString = json.dumps(data)
with open('Conv2DTest.json', 'w') as f:
    f.write(jsonString)
print('TestBench generate finished.')